import torch
import torch.nn.functional as F
from copy import deepcopy
torch.manual_seed(0)

bn = torch.nn.BatchNorm2d(2)
bn2 = torch.nn.BatchNorm2d(2)

rand_y = torch.randn(10, 2, 32, 32)
rand_x1 = torch.randn(10, 2, 32, 32) * -10
rand_x2 = torch.randn(10, 2, 32, 32) * 10

opt = torch.optim.Adam(bn.parameters(), lr=0.01)
bn.train()

print("fw0", bn.running_mean)
hat_x = bn(rand_x1)
print("fw1", bn.running_mean)
hat_x2 = bn(rand_x2)
print("fw2", bn.running_mean)

print("case 2")
print("fw0", bn2.running_mean)
hat_x = bn2(rand_x1)
bn2.track_running_stats = False
print("fw1", bn2.running_mean)
hat_x2 = bn2(rand_x2)
print("fw2", bn2.running_mean)
